import logging

import cv2
import os
import ffmpeg
from tqdm import tqdm


class VideoProcessor:

    def __init__(self, video_path, temp_path, frame_count=50):

        self.video_path = video_path  # 读取视频路径
        self.temp_path = temp_path  # 临时文件路径
        self.frame_count = frame_count  # 一次读取帧数

        if not os.path.exists(self.temp_path):
            os.makedirs(self.temp_path)

        self.fps = None

    @staticmethod
    def process_frame(frame):
        """
        图像调整 将每一帧 调整为 320x180 这是为迎合model推理
        """
        resized_frame = cv2.resize(frame, (320, 180), interpolation=cv2.INTER_CUBIC)
        return resized_frame

    def write_frames(self, frame_list, outfile):
        """
        将frame_list 的帧 写入一个mp4 文件
        """

        height, width = frame_list[0].shape[:2]
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')

        out = cv2.VideoWriter(outfile, fourcc, self.fps, (width, height))
        for frame in frame_list:
            out.write(frame)
        out.release()

    def read_frames(self, count):
        """
        :param count: 每个视频写入帧数
        :return:
        """
        cap = cv2.VideoCapture(self.video_path)
        self.fps = cap.get(cv2.CAP_PROP_FPS)  # 设置视频fps

        frame_list = []
        num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        logging.info(f"num_frames : {num_frames}")

        pbar = tqdm(total=num_frames)  # 进度条
        write_num = 0
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # 视频帧列表full 写视频
            if len(frame_list) % count == 0 and frame_list:
                out_name = f"{write_num + 1}.mp4"
                write_num += 1
                outfile = os.path.join(self.temp_path, out_name)
                self.write_frames(frame_list, outfile)
                frame_list.clear()

            frame_list.append(VideoProcessor.process_frame(frame))
            pbar.update(1)

        # 跳出循环后将剩余帧写入一个视频
        if frame_list:
            out_name = f"{write_num + 1}.mp4"
            write_num += 1
            outfile = os.path.join(self.temp_path, out_name)
            self.write_frames(frame_list, outfile)

        cap.release()
        pbar.close()

    def merge_videos(self, output_name):
        """
        合并产生的临时视频
        :param output_name:
        :return:
        """
        video_files = sorted([os.path.join(self.temp_path, f) for f in os.listdir(self.temp_path)])
        assert all([os.path.isfile(f) for f in video_files]), "Some input files are not files"
        assert all([vf.endswith('.mp4') for vf in video_files]), "Only support *.mp4 file format"

        input_args = []
        for vf in video_files:
            input_args += [ffmpeg.input(vf)]
        output_args = ffmpeg.output(ffmpeg.concat(*input_args), output_name, y="-y")

        # 合并视频
        output_args.run()

    def processing(self):

        self.read_frames(count=self.frame_count)
        self.merge_videos('temp.mp4')

        # 删除临时文件
        for file in os.listdir(self.temp_path):
            temp_file = os.path.join(self.temp_path, file)
            if os.path.exists(temp_file):
                os.remove(temp_file)

        logging.info('video process finished')

if __name__ == '__main__':
    video_processor = VideoProcessor(video_path='path/to/video.mp4')
    video_processor.processing()
